#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/10/24 6:54
@Author   : liujingmao
@File     : 1.XMLAgent示例.py
"""
import dotenv
from langchain.agents import create_xml_agent, AgentExecutor
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
# from langchain_core.pydantic_v1 import BaseModel, Field
from pydantic import BaseModel, Field

dotenv.load_dotenv()


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")


# 1.定义工具与工具列表
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答有关时事的问题时，可以调用该工具。"
        "该工具的输入是搜索查询语句。"
    ),
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)
dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DallEArgsSchema,
)
tools = [google_serper, dalle]

# 2.定义工具调用agent提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("human", """You are a helpful assistant. Help the user answer any questions.

You have access to the following tools:

{tools}

In order to use a tool, you can use <tool></tool> and <tool_input></tool_input> tags. You will then get back a response in the form <observation></observation>
For example, if you have a tool called 'search' that could run a google search, in order to search for the weather in SF you would respond:

<tool>search</tool><tool_input>weather in SF</tool_input>
<observation>64 degrees</observation>

When you are done, respond with a final answer between <final_answer></final_answer>. For example:

<final_answer>The weather in SF is 64 degrees</final_answer>

Begin!

Previous Conversation:
{chat_history}

Question: {input}
{agent_scratchpad}"""),
])

# 3.创建大语言模型
llm = ChatOpenAI(model="gpt-4o-mini")

# 4.创建agent与agent执行者
agent = create_xml_agent(
    prompt=prompt,
    llm=llm,
    tools=tools,
)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print(agent_executor.invoke({"input": "刀郎下次的演唱会的时间是什么时候？", "chat_history": ""}))

'''
Entering new AgentExecutor chain...
<tool>google_serper</tool><tool_input>马拉松 世界记录 2023当地时间24日，在2023年柏林马拉松赛上，肯
尼亚长跑名将埃鲁德·基普乔格以2小时02分41秒的成绩夺冠。 
这是基普乔格赢得的第五个柏林马拉松赛冠军。 基普乔格在2022年的柏林马拉松赛上以2小时01分09秒的成绩创造世界纪录。<final_answer>马拉松的世界记录是2小时01分09秒，由肯尼亚选手埃鲁德·基普乔格于2022年柏林马拉松赛上创造。</final_answer>

> Finished chain.
{'input': '马拉松的世界记录是多少？', 'chat_history': '', 'output': '马拉松的世界记录是2小时01分09秒，由肯尼亚选手埃鲁德·基普乔格于2022年柏林马拉松赛上创造。'}
'''

'''
 Entering new AgentExecutor chain...
<tool>google_serper</tool><tool_input>next Olympic Games dateFri, Feb 6, 2026 – Sun, Feb 22, 2026<final_answer>下次举办奥运会的时间是2026年2月6日至2月22日。</final_answer>

> Finished chain.
{'input': '下次举办奥运会的时间是？', 'chat_history': '', 'output': '下次举办奥运会的时间是2026年2月6日至2月22日。'}

Process finished with exit code 0
'''

## 下面的答案显然是不对的

'''
 Entering new AgentExecutor chain...
<tool>google_serper</tool><tool_input>刀郎 演唱会 时间 20238月15日，刀郎再次宣布“复出”，他通过个人社交平台宣布将于8月30日晚举行线上演唱会。 他在视频中讲述了线上演唱会的内容，他说：“我走了很长一段时间， ... 2024刀郎北京演唱会演出时间是2024年12月28日-29日和2025年1月4日-5日，将在北京华熙LIVE·五棵松·场馆举办。 刀郎演唱会2024巡演演出安排 · 刀郎澳门演唱会: 2024-11-08/09 20:00 · 刀郎厦门演唱会: 2024.11.30-12.01 · 刀郎合肥演唱会 · 刀郎北京演唱会 · 刀郎成都演唱会 · 刀郎广州演唱 ... 近日，刀郎宣布“复出”，他通过个人社交平台宣布将于8月30日晚举行线上演唱会。还有消息称，今年，刀郎还将在全国多地开巡回演唱会。 刀郎“复出”. 刀郎全国巡演厦门站具体信息！ 演出信息. 演出时间. 11月30日19:00. 12月1日19:00. 演出地点. 厦门奥林匹克体育中心 ... 他是满腹才华的音乐达人，他是性情坦荡的全能艺人，他就是刀郎。刀郎2024演唱会时间排期表公布，刀郎2024演唱会成都站将于9月21-22日举办，一起来现场聆听 ... 刀郎全国巡演厦门站具体信息！ 演出信息. 演出时间. 11月30日19:00. 12月1日19:00. 演出地点. 厦门奥林匹克体育中心 ... 00:00 通过个人社交平台宣布将于8月30日晚举行线上演唱会。截止19日下午3时，"刀郎《山歌响起的地方》巡回演唱会"已在成都、广州通过审批或许可。 山歌想起的地方·刀郎2024巡回演唱会合肥站，2024年12月14日-12月15日在合肥少荃体育中心体育馆举行，具体演出内容见正文。 时间：2024.09.21-09.22地点：成都凤凰山体育公园综合体育馆门票价格：480680880,1080,1280刀郎，本名罗林，华语流行乐男歌手、词曲作者、音乐人， ...<final_answer>刀郎下次的演唱会时间是2024年12月28日-29日和2025年1月4日-5日，将在北京华熙LIVE·五棵松·场馆举办。此外，他还将在2024年11月30日和12月1日举办厦门演唱会。</final_answer>

> Finished chain.
{'input': '刀郎下次的演唱会的时间是什么时候？', 'chat_history': '', 'output': '刀郎下次的演唱会时间是2024年12月28日-29日和2025年1月4日-5日，将在北京华熙LIVE·五棵松·场馆举办。此外，他还将在2024年11月30日和12月1日举办厦门演唱会。'}

Process finished with exit code 0

'''
